"""


"""
import os
import shutil
import argparse
from PIL import Image

image_extensions = {'.jpg', '.jpeg', '.JPG', '.JPEG'}


def parse_args():
    parser = argparse.ArgumentParser(
        description="Put all places 365 images in single folder.")
    parser.add_argument(
        "-i",
        "--input_dir",
        required=True,
        type=str,
        help="input folder: the folder containing unzipped places 365 files")
    args = parser.parse_args()
    return args


def resize_from_path(p, new_p):
    a = Image.open(p)
    if a.getbands() != ('R', 'G', 'B'):
        os.remove(p)
        return
    # todo: resize 512 to larger size
    a.resize((256, 256), Image.BICUBIC).save(new_p)
    return

def genlist(image_dir):
    image_list = []
    for filename in os.listdir(image_dir):
        path = os.path.join(image_dir, filename)
        if os.path.isdir(path):
            image_list = image_list + genlist(path)
        else:
            ext = os.path.splitext(filename)[1]
            if ext in image_extensions:
                image_list.append(os.path.join(image_dir, filename))
    return image_list


args = parse_args()
output_dir = args.input_dir + '_colorize_train'
os.makedirs(output_dir, exist_ok=True)

flist = genlist(args.input_dir)

for i, p in enumerate(flist):
    if os.path.getsize(p) != 0:
        # resize it and then rename it
        target_p = os.path.join(output_dir, str(i) + '.jpg')
        resize_from_path(p, target_p)
print('done')